Support LLaDa-Image by lucasruan1618 · Pull Request #14815 · huggingface/diffusers

lucasruan1618 · 2026-09-20T03:37:50Z

What does this PR do?

This PR adds native LLaDA-Image support to Diffusers. It introduces LLaDAImagePipeline, a single pipeline for text-to-image generation, VQ-conditioned generation, and instruction-guided image editing.

The implementation ports the published LLaDA-Image inference architecture while using Diffusers' standard component registration, serialization, device placement, CPU offloading, group offloading, attention processor, and pipeline loading interfaces.

Pipeline architecture

LLaDAImagePipeline registers the eight components already described by the published model_index.json:

Component	Role in inference
`text_encoder` and `tokenizer`	The custom LLaDA2 language model produces prompt representations and, in VQ mode, autoregressively generates image-token IDs.
`queryformer`	Appends learned image-generation queries to the tokenized prompt before the LLaDA2 backbone runs.
`text_projection`	Maps the LLaDA2 hidden states into the caption-feature dimension expected by the denoising transformer.
`sigvq`	Encodes a reference image into semantic VQ features for editing, or maps MLLM-generated VQ IDs into equivalent features for VQ-conditioned generation.
`transformer`	Runs flow-matching denoising over Flux2 latent patches, conditioned on text and optional SigVQ features/source latents.
`vae`	Encodes the source image for editing and decodes final Flux2 latents into images.
`scheduler`	Provides the published flow-matching timestep schedule.

The pipeline is loaded with the usual DiffusionPipeline.from_pretrained path. The published LLaDA2 text encoder is custom remote code, so users must pass trust_remote_code=True when loading the official checkpoint.

import torch

from diffusers import LLaDAImagePipeline


pipe = LLaDAImagePipeline.from_pretrained(
    "inclusionAI/LLaDA-Image",
    dtype=torch.bfloat16,
    trust_remote_code=True,
)
pipe.enable_model_cpu_offload()

text_encoder remains resident because the pipeline directly calls its embedding layer and language backbone; the QueryFormer output must be inserted between those calls. The remaining model components participate in the normal offloading sequence: QueryFormer, text projection, SigVQ, denoising transformer, and VAE.

Supported inference modes

Text-to-image

generation_mode="text" is the default path. The pipeline encodes the positive prompt and, when guidance_scale > 1, an empty or supplied negative prompt. It then denoises random Flux2 latent patches using classifier-free guidance.

image = pipe(
    prompt="A red fox walking through fresh snow, cinematic photography",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]

VQ-conditioned generation

generation_mode="vq" asks the LLaDA2 image-generation head for VQ token IDs. The pipeline converts those IDs to SigVQ semantic features and supplies them to the denoising transformer alongside prompt features. The frontend VQ grid is capped at 512 pixels on its longest side, matching the reference implementation.

image = pipe(
    prompt="A friendly robot tending a rooftop garden, colorful editorial illustration",
    generation_mode="vq",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
).images[0]

Image editing

generation_mode="editing" requires an input image. The pipeline normalizes and encodes that image twice: SigVQ produces semantic image features, and the Flux2 VAE produces source latents. The denoising transformer receives both forms of conditioning with the text instruction.

from diffusers.utils import load_image


source = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"
)
image = pipe(
    prompt="Turn it into a watercolor painting with a soft blue background",
    image=source,
    generation_mode="editing",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=5.0,
).images[0]

The pipeline validates mode-specific inputs before inference. Text and VQ modes reject image; editing requires it; VQ dimensions must be divisible by 16; and all output dimensions must match the Flux2 VAE and latent-patch scaling requirements.

New public model components

The PR adds and exports four serializable Diffusers models:

LLaDAImageTransformer2DModel, the variable-resolution denoising transformer.
LLaDAImageQueryFormerModel, which refines learned generation queries against token embeddings.
LLaDAImageTextProjectionModel, which connects LLaDA2 hidden states to transformer caption features.
LLaDAImageSigVQModel, which supports both image-to-VQ encoding and VQ-token-to-semantic-feature lookup.

The transformer preserves the reference model's list-valued output so a batch may contain samples with different spatial shapes. Its RoPE cache is bypassed only while torch.compile traces the model, avoiding module-state mutation during export while retaining the normal eager-mode cache.

Diffusers integration details

Adds lazy imports and dependency dummy objects for the pipeline and all four models.
Uses the standard pipeline loader without a checkpoint-specific from_pretrained override.
Extends the generic custom-component downloader to include sibling Python files, allowing custom components such as LLaDA2 to import their local configuration and fused-MoE modules.
Restores the standard default RoPE registry entry required by the checkpoint's custom text encoder when Transformers 5 omits it.
Re-materializes LLaDA2's non-persistent RoPE frequency buffer after Transformers 5 meta-device loading. Without this repair, the buffer contains invalid values and VQ sampling emits text/control tokens outside the SigVQ codebook.
Keeps the published model-index component names and configuration fields compatible with the official checkpoint.
Adds model and pipeline API documentation, plus entries in the documentation table of contents.
Supports save_pretrained / from_pretrained, dtype loading, device maps, CPU/disk offload, model CPU offload, group offload, callbacks, batching, and supported image output types.

Tests

tests/models/transformers/test_models_transformer_llada_image.py: model serialization, deterministic outputs, dtype loading, CPU/disk/group offload, gradient checkpointing, attention processor behavior, compilation, and direct QueryFormer, projection, and SigVQ forward coverage.
tests/pipelines/llada_image/test_pipeline_llada_image.py: shared pipeline contracts for loading, batching, callbacks, serialization, dtype handling, accelerator integration, and offloading; plus focused VQ and image-editing tests.

Validation performed:

pytest tests/models/transformers/test_models_transformer_llada_image.py -q: 43 passed, 12 skipped.
pytest tests/pipelines/llada_image/test_pipeline_llada_image.py -q: 38 passed, 1 skipped.
make quality, utils/check_dummies.py, utils/check_copies.py, and git diff --check: passed.
Manual full-checkpoint runs on an NVIDIA A100 80 GB, bfloat16, model CPU offload, and 50 steps:
- Text-to-image at 1024×1024, CFG 5.0, seed 42: llada_image_text_to_image.png.
- VQ-conditioned generation at 512×512, CFG 5.0, seed 43, using Transformers 5.15.1: llada_image_vq_transformers_5.png.
- Image editing at 512×512, CFG 5.0, seed 44, using the generated fox image as the source: llada_image_edit.png.

The skips cover generic test utilities that cannot operate on the transformer's intentionally list-valued input/output interface, plus AOT package loading, which currently cannot deserialize list-valued inputs. The standard eager, dynamic-shape, and repeated-block torch.compile tests pass.

Self-review

Verdict: READY

No blocking or non-blocking issues remain. The Transformers 5.15.1 VQ failure was traced to the official remote text encoder's non-persistent RoPE buffer being initialized on the meta device during sharded loading. The corrected buffer reproduces the Transformers 4.57.6 reference tokens exactly in the reduced comparison, and the full 512×512 Transformers 5 output is byte-identical to the reference-runtime output.

No likely-dead inference paths were found. The text, VQ, and editing paths are all traced from LLaDAImagePipeline.__call__ and covered by tests. The official configuration schema was checked against every new model constructor, and the port preserves upstream model math apart from the compile-safe RoPE cache guard and Diffusers device/offloading integration.

Before submitting

Who can review?

Pipelines and models: @yiyixuxu and @asomoza

github-actions · 2026-09-20T07:36:33Z

Hi @lucasruan1618, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice.

Once the PR links an issue (or gets the no-issue-needed label), you can ignore this message — it stays here as a comment, but it no longer applies.

sergereview

🤗 Serge says:

Large, well-structured port of LLaDA-Image: four models in transformer_llada_image.py, a pipeline with an output dataclass, docs, dummies and both test layers. The model file follows the Z-Image/Lumina lineage closely (ADALN constant, per-token modulation, pad_sequence batching) and the attention classes correctly use AttentionModuleMixin + dispatch_attention_fn. A few things need attention before merge.

Correctness

Global monkeypatch of a Transformers registry at import time. pipeline_llada_image.py mutates transformers.modeling_rope_utils.ROPE_INIT_FUNCTIONS at module import, so importing diffusers.pipelines.llada_image silently changes behaviour for every Transformers model in the process that resolves rope_type="default".
Random weight initialization in a load-only model. LLaDAImageTransformer2DModel.__init__ and LLaDAImageQueryAttention.__init__ call nn.init.normal_ / xavier_uniform_ / zeros_. For checkpoint-loaded inference these are pure overhead and can mask a missing key (a layer that fails to load gets plausible random init instead of the usual meta/empty tensor signal). Drop them or move them behind an explicit init hook.
t_embedder timestep embedding is computed in the caller's autocast context, unlike the reference TimestepEmbedder in transformer_z_image.py, which wraps the frequency computation in torch.amp.autocast(..., enabled=False). Confirm numerics under bf16 autocast.
Undocumented LLaDAImageSigVQModel config args. No Args: section at all while sibling models document each config value; [[autodoc]] renders an empty parameter list for a model with 13 config knobs.
generate_vq_tokens depends on remote-code API and magic constants. self.text_encoder.generate_bd_image_logic(...), image_token_offset = 157184, <|reserved_token_N|> names and block_length=32, steps=8, cfg_scale=2.0 are hard-coded against one remote checkpoint with no hasattr guard — a different text encoder gives an AttributeError deep in the call. Add an up-front check plus a comment naming where the offset comes from.

Docs

Both new pages (api/models/llada_image_transformer2d.md, api/pipelines/llada_image.md) omit the Apache license header that every other page in docs/source/en/api/pipelines/ carries (compare longcat_image.md).
docs/source/en/_toctree.yml titles the model entry LLaDA-Image while surrounding entries use the class name (LatteTransformer3DModel, LongCatImageTransformer2DModel).

Unrelated change

The pipeline_utils.py download() hunk broadens custom-component allow patterns from {component}/{module}.py to {component}/*.py — a behaviour change for every pipeline with custom components. It belongs in its own PR with its own test, or at minimum needs justification in the description (currently unmentioned).

Tests

Good coverage overall (model + pipeline layers, memory mixin, no slow/LoRA tests, per project convention). Concerns:
- test_transformers_5_default_rope_is_materialized hand-builds a SimpleNamespace config and a bare nn.Module rotary_emb, then asserts the pipeline wrote the values it itself computed — verifying the pipeline against itself rather than the real Transformers rotary-embedding contract (see testing.md on mocks / call-level test doubles).
- test_vq_conditioned monkeypatches generate_bd_image_logic onto the text encoder, so the VQ path is exercised only against a stub returning exactly what the pipeline expects. Given the magic 157184 offset and the token-count validation, that is the part most worth testing against real behaviour.
- Six @pytest.mark.skips on the training/compile mixins are all justified by the list-valued input/output signature, which is fine — but the number of shared tests opted out of is worth a maintainer look, since it stems from forward taking list[torch.Tensor] rather than a batched tensor.

serge v0.1.0 · model: claude-opus-5 · 21 LLM turns · 27 tool calls · 731.0s · 2153181 in / 57702 out tokens

sergereview · 2026-09-25T16:01:26Z

+
+# The official LLaDA2 remote model uses this Transformers 4 compatibility entry. Transformers 5 no longer
+# registers it, so make it available before DiffusionPipeline loads the custom text encoder.
+if "default" not in ROPE_INIT_FUNCTIONS:


This mutates a global Transformers registry at import time; since diffusers/pipelines/__init__.py lazily exposes llada_image, any import path changes RoPE initialization process-wide for every Transformers model resolving rope_type="default", not just the LLaDA2 encoder loaded here. Please scope it: register the entry inside __init__ (or just before the remote model loads) and restore it afterwards, or pass the init function directly to the rotary module materialized below — you already hold the _default_rope_parameters reference, so line 117 does not need the registry at all.

sergereview · 2026-09-25T16:01:26Z

+            self.text_encoder.to(execution_device)
+
+        prompts = [prompt] if isinstance(prompt, str) else prompt
+        image_token_offset = 157184


image_token_offset = 157184 is an unexplained magic constant tied to one checkpoint's tokenizer, and generate_bd_image_logic (line 283) is a remote-code method only the official LLaDA2 encoder exposes — any other PreTrainedModel gives an AttributeError mid-loop. Please (a) comment where 157184 comes from (e.g. the id of the first image VQ token in the LLaDA2 vocabulary) so it can be re-derived, and (b) fail early with a clear message, e.g.

if not hasattr(self.text_encoder, "generate_bd_image_logic"): raise ValueError( "`generation_mode='vq'` requires the LLaDA2 text encoder, which exposes `generate_bd_image_logic`." )

Same for hard-coded block_length=32, steps=8, cfg_scale=2.0 — surface them as arguments or note they are the published reference defaults.

sergereview · 2026-09-25T16:01:26Z

+            nn.Linear(semantic_feat_dim, dim, bias=True),
+        )
+
+        nn.init.normal_(self.semantic_embedder[1].weight, mean=0.0, std=0.02)


These four nn.init calls (and nn.init.normal_(self.sigvq_pad_token, ...) on line 430) run on every instantiation including from_pretrained, where the checkpoint immediately overwrites them; beyond wasted work on a 3840-dim model, a random init masks a missing/renamed key as plausible-looking noise instead of an obviously uninitialized tensor. The comparable Z-Image transformer does not init in __init__ — recommend dropping these and the xavier_uniform_/zeros_ pair in LLaDAImageQueryAttention.__init__ (lines 1248-1249), unless training-from-scratch parity requires them.

sergereview · 2026-09-25T16:01:26Z

+
+    def forward(self, timestep: torch.Tensor, hidden_dtype: torch.dtype) -> torch.Tensor:
+        half_dim = self.frequency_embedding_dim // 2
+        frequencies = torch.exp(


The reference TimestepEmbedder this is ported from (transformer_z_image.py, timestep_embedding) computes the sinusoidal frequencies inside torch.amp.autocast(..., enabled=False). Here cos/sin are float32 explicitly but the surrounding autocast is not disabled, so under bf16 autocast the self.mlp(...) call and downstream fusion can differ from the reference. Confirm the numerics match the original under bf16, or mirror the explicit autocast-disable.

sergereview · 2026-09-25T16:01:26Z

+
+
+class LLaDAImageSigVQModel(ModelMixin, ConfigMixin, AttentionMixin):
+    r"""


Unlike LLaDAImageTransformer2DModel, LLaDAImageQueryFormerModel and LLaDAImageTextProjectionModel, this docstring has no Args: section, so the [[autodoc]] LLaDAImageSigVQModel block in docs/source/en/api/models/llada_image_transformer2d.md renders 13 config parameters with no documentation. Please document image_size, patch_size, codebook_size, codebook_embed_dim, semantic_embed_dim and friends in the sibling classes' style.

sergereview · 2026-09-25T16:01:26Z

-            # add custom component files
-            allow_patterns += [f"{k}/{f}.py" for k, f in custom_components.items()]
+            # Add custom component modules and their local Python dependencies.
+            allow_patterns += [f"{folder_name}/*.py" for folder_name in custom_components]


This widens the download allow-list for all pipelines with custom components from {component}/{module}.py to {component}/*.py, changing what gets fetched for every existing custom-component repo, not just LLaDA-Image — and the PR description does not mention it. It also discards the mapping's values: _get_custom_components_and_folders returns {component: module_name} precisely so only the declared module is allowed, and pulling every .py in the folder is a broader trust surface for remote-code repos. Please split this into its own PR with a test covering a custom component whose module imports a local sibling, or narrow it to the specific dependency-resolution case needed here.

sergereview · 2026-09-25T16:01:26Z

@@ -0,0 +1,57 @@
+# LLaDA-Image


Missing the Apache license header comment that every other page under docs/source/en/api/pipelines/ starts with (see longcat_image.md). Same for docs/source/en/api/models/llada_image_transformer2d.md.

Suggested change

# LLaDA-Image



# LLaDA-Image

sergereview · 2026-09-25T16:01:26Z

      - local: api/models/latte_transformer3d
        title: LatteTransformer3DModel
+      - local: api/models/llada_image_transformer2d
+        title: LLaDA-Image


Entries in the models section are titled by class name (LatteTransformer3DModel, LongCatImageTransformer2DModel, Krea2Transformer2DModel). LLaDA-Image here also collides with the identically titled pipeline entry at line 623, making the two hard to tell apart in the sidebar.

Suggested change

title: LLaDA-Image

title: LLaDAImageTransformer2DModel

sergereview · 2026-09-25T16:01:26Z

+            pytest.skip("This regression only affects Transformers 5 and later.")
+
+        components = self.get_dummy_components()
+        rotary_emb = torch.nn.Module()


Both sides of the assertion come from the code under test: a torch.nn.Module with a SimpleNamespace config stands in for the real rotary embedding, and the test asserts the pipeline wrote back exactly what _default_rope_parameters computes from that same namespace. It stays green if Transformers renames original_inv_freq, changes attention_scaling semantics, or stops using rope_theta — precisely the contract this workaround depends on. Per testing.md, exercise the real component: instantiate an actual Transformers rotary module (or the real text encoder under the version guard) and assert its inv_freq is no longer a meta/zero tensor after pipeline construction.

sergereview · 2026-09-25T16:01:26Z

+    def test_vq_conditioned(self):
+        pipe = self.get_pipeline().to(torch_device)
+
+        def generate_bd_image_logic(text_encoder, data, block_length, steps, gen_length, cfg_scale):


This stub returns input_ids concatenated with arange(gen_length) + 157184, i.e. exactly the layout generate_vq_tokens slices and offsets, so the test cannot fail on a wrong offset, a wrong slice window, or a mismatched token count — the three things the validation at pipeline_llada_image.py:295-299 exists to catch. At minimum add cases driving those error paths (a stub returning too few tokens, and one returning an out-of-codebook id) so the ValueErrors are covered.

lucasruan1618 and others added 3 commits September 6, 2026 14:20

Add Llada Image pipeline

bbca6bd

Fix vq mode for transformers 5.

3b8a581

Merge branch 'huggingface:main' into add-llada-image

9f87f34

github-actions Bot added documentation Improvements or additions to documentation models tests utils pipelines size/L PR with diff > 200 LOC labels Sep 20, 2026

sayakpaul requested a review from kashif September 25, 2026 07:56

sayakpaul added the no-issue-needed for PRs that do not require link to an issue label Sep 25, 2026

sayakpaul requested review from zucchini-nlp and removed request for zucchini-nlp September 25, 2026 07:57

sergereview Bot reviewed Sep 25, 2026

View reviewed changes

Navigation Menu

Support LLaDa-Image - #14815

Support LLaDa-Image#14815
lucasruan1618 wants to merge 3 commits into
huggingface:mainfrom
lucasruan1618:add-llada-image

lucasruan1618 commented Sep 20, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Sep 20, 2026 •

edited

Loading

Uh oh!

sergereview Bot left a comment

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

sergereview Bot Sep 25, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		class LLaDAImageSigVQModel(ModelMixin, ConfigMixin, AttentionMixin):
		r"""

-# LLaDA-Image
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+# LLaDA-Image

Conversation

lucasruan1618 commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Pipeline architecture

Supported inference modes

Text-to-image

VQ-conditioned generation

Image editing

New public model components

Diffusers integration details

Tests

Self-review

Before submitting

Who can review?

Uh oh!

github-actions Bot commented Sep 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sergereview Bot left a comment

Choose a reason for hiding this comment

Correctness

Docs

Unrelated change

Tests

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

lucasruan1618 commented Sep 20, 2026 •

edited

Loading

github-actions Bot commented Sep 20, 2026 •

edited

Loading